#!/usr/bin/python3

import rospy
import math
import tf
import tf2_ros
from geometry_msgs.msg import Twist, TransformStamped, PointStamped, PoseStamped, PolygonStamped
import time
import numpy as np
from std_msgs.msg import Float64MultiArray, Header, Bool
from sensor_msgs.msg import LaserScan, JointState
from gazebo_msgs.msg import ModelStates
from visualization_msgs.msg import Marker, MarkerArray
import copy
from tf2_geometry_msgs import tf2_geometry_msgs
import sys
from pogo_agv.srv import GetPoseFromPoints, GetPoseFromPointsRequest, GetPoseFromPointsResponse
from pogo_agv.msg import MovePoseWithSupportPoints
from data_fitting import PointMatcher2
import threading

class pogo_agv_service():
    def __init__(self,nodename, ns_prefix):
        rospy.init_node(nodename)
        # debug_prefix = "/" + ns_prefix + "/"
        debug_prefix = ""
        self.debug_prefix = debug_prefix

        self.tfBuffer = tf2_ros.Buffer()
        self.tfListener = tf2_ros.TransformListener(self.tfBuffer)
        self.tfBroadcaster=tf2_ros.TransformBroadcaster()

        self.real_base_link = ns_prefix + "/target_moving_pose"
        # self.real_support_link = ns_prefix + "/target_support_pose"
        self.base_link = ns_prefix + "/base_link"

        self.real_agv_hight = None
        self.real_agv_pose = None
        self.real_support_points = None
        # self.is_support_point_mode = False

        # self.sim_agv_pose = rospy.Publisher(debug_prefix + "real/sim_agv_pose", PoseStamped, queue_size = 10)
        self.sim_cmd_vel = rospy.Publisher(debug_prefix + "cmd_vel", Twist,queue_size=10)
        self.sub_move_pose = rospy.Subscriber(debug_prefix + "move_to_pose", PoseStamped, self.move_pose_cb)
        # self.sub_move_pose = rospy.Subscriber(debug_prefix + "move_to_support_points", PolygonStamped, self.move_pose_support_points)
        self.sub_move_pose = rospy.Subscriber(debug_prefix + "move_to_support_points", MovePoseWithSupportPoints, self.move_pose_support_points)

        self.agv_joint_pub = rospy.Publisher(debug_prefix + 'agv_joint_states', JointState, queue_size=1)
        self.srv_get_support_pose = rospy.Service(debug_prefix + "get_support_pose", GetPoseFromPoints, self.get_support_pose)
        self.agv_support_names = [ns_prefix + "/link3_" + str(i) for i in [4,3,1,2]]
        self.agv_lift_joints = ["joint3_" + str(i) for i in [4,3,1,2]]

        self.joint_name_index = None
        self.realtime_tracking_switch_on = True
        self.joint_states={}

        self.sim_ptp_controller_command = rospy.Publisher(debug_prefix + 'ptp_controller/command', Float64MultiArray, queue_size=10)
        self.real_joint_states = rospy.Subscriber(debug_prefix + "cmd_joint_states",JointState,self.realJointStates_cb)
        self.turn_on_realtime_tracking = rospy.Subscriber(debug_prefix + "turn_on_realtime_tracking",Bool, self.turnOnRealtimeTracking_cb)
        self.joint_states_topic = rospy.Subscriber(debug_prefix+"joint_states",JointState,self.joint_state_cb)

        self.lock = threading.Lock()
        self.reset_agv_pose_control()
        # self.js = JointState()
        # self.js.header = Header()
        # self.js.name = []
        # self.js.position = []
        # self.js.velocity = []
        # self.js.effort = []

    def joint_state_cb(self,msg):
        for i in range(len(msg.name)):
            self.joint_states[msg.name[i]]=msg.position[i]

    def realJointStates_cb(self,msg):
        if(not self.realtime_tracking_switch_on):
            return 

        if(self.joint_name_index is None):
            para_name = self.debug_prefix + "ptp_controller/joints"
            if(rospy.has_param(para_name)):
                self.joint_names = rospy.get_param(para_name)
                self.joint_name_index = {self.joint_names[i]:i for i in range(len(self.joint_names))}
            else:
                return

        sent_command=[0.0]*len(self.joint_names)
        for i in range(len(msg.name)):
            sent_command[self.joint_name_index[msg.name[i]]] = msg.position[i]
        
        self.sim_ptp_controller_command.publish(Float64MultiArray(data=sent_command))
        return

    def turnOnRealtimeTracking_cb(self,msg):
        print(msg)
        return

    def get_support_pose(self,req):

        # return GetPoseFromPointsResponse(pose=TransformStamped())

        target_points = []
        for i in req.points:
            sp = [i.point.x, i.point.y, i.point.z]
            target_points.append(sp)
        
        support_points = []
        for i in self.agv_support_names:
            pos = self.tfBuffer.lookup_transform(self.base_link, i, rospy.Time(0))
            sp = [pos.transform.translation.x, pos.transform.translation.y, pos.transform.translation.z]
            support_points.append(sp)

        matcher = PointMatcher2(np.array(target_points),np.array(support_points))
        result = matcher.match_points()
        trans = PoseStamped()

        trans.header.frame_id = self.base_link
        trans.header.stamp = rospy.Time.now()

        trans.pose.position.x = result[1][0]
        trans.pose.position.y = result[1][1]
        trans.pose.position.z = result[1][2]

        rot = np.eye(4,4)
        rot[:3,:3]=result[0]
        quat = tf.transformations.quaternion_from_matrix(rot)
        trans.pose.orientation.x = quat[0]
        trans.pose.orientation.y = quat[1]
        trans.pose.orientation.z = quat[2]
        trans.pose.orientation.w = quat[3]

        return GetPoseFromPointsResponse(pose = trans)

    def move_pose_support_points(self,msg):
        self.lock.acquire()
        # self.reset_agv_pose_control()
        self.real_support_points = copy.deepcopy(msg)
        self.support_points_control_step()
        # self.real_agv_pose = None
        self.lock.release()

    def move_pose_cb(self,msg):
        self.lock.acquire()
        # self.reset_agv_pose_control()
        self.real_agv_pose = copy.deepcopy(msg)
        # self.real_support_points = None
        self.lock.release()

    def support_points_control_step(self):
        # print("support points control:",self.real_support_points)
        # return

        #firstly, get the zero point of the pogos on agv
        support_points = []
        for i in self.agv_support_names:
            pos = self.tfBuffer.lookup_transform(self.base_link, i, rospy.Time(0))
            sp = [pos.transform.translation.x, pos.transform.translation.y, pos.transform.translation.z]
            support_points.append(sp)

        zero_points = []
        for i in range(len(self.agv_lift_joints)):
            zp = support_points[i][2]-self.joint_states[self.agv_lift_joints[i]]
            zero_points.append(zp)

        self.real_support_points.pose.header.stamp = rospy.Time(0)
        pose_in_agv = self.tfBuffer.transform(self.real_support_points.pose,self.base_link)
        #second, convert the pose to rpy and z
        quat = [pose_in_agv.pose.orientation.x,pose_in_agv.pose.orientation.y,pose_in_agv.pose.orientation.z, pose_in_agv.pose.orientation.w]
        rpy = tf.transformations.euler_from_quaternion(quat)
        trans = tf.transformations.compose_matrix(angles=[rpy[0],rpy[1],0],translate=[0,0,pose_in_agv.pose.position.z])
        
        sent_command = JointState()
        target_pogo_points=[]
        target_points = []
        for i in range(len(self.real_support_points.points)):
            pt = self.real_support_points.points[i]
            sp = np.array([pt.point.x, pt.point.y, pt.point.z,1.0])
            points_in_agv = trans@sp

            sent_command.name.append(self.agv_lift_joints[i])
            sent_command.position.append(points_in_agv[2]-zero_points[i])

            target_pogo_points.append([support_points[i][0],support_points[i][1],points_in_agv[2]])
            target_points.append([pt.point.x,pt.point.y,pt.point.z])

        matcher = PointMatcher2(np.array(target_pogo_points),np.array(target_points))
        result = matcher.match_points()

        rot = np.eye(4,4)
        rot[:3,:3]=result[0]
        rot[:3,3]=result[1]
        # rot1 = tf.transformations.inverse_matrix(rot)

        cp = self.real_support_points.pose.pose
        rot2 = tf.transformations.translation_matrix([cp.position.x,cp.position.y, cp.position.z])@tf.transformations.quaternion_matrix([
            cp.orientation.x, cp.orientation.y, cp.orientation.z, cp.orientation.w
        ])
        rot3 = rot2@rot
        
        agv_pose=PoseStamped()
        agv_pose.header.frame_id = self.real_support_points.pose.header.frame_id
        agv_pose.header.stamp=rospy.Time.now()
        agv_pose.pose.position.x = rot3[0,3]
        agv_pose.pose.position.y = rot3[1,3]
        agv_pose.pose.position.z = rot3[2,3]

        quat = tf.transformations.quaternion_from_matrix(rot3)
        agv_pose.pose.orientation.x = quat[0]
        agv_pose.pose.orientation.y = quat[1]
        agv_pose.pose.orientation.z = quat[2]
        agv_pose.pose.orientation.w = quat[3]

        self.real_agv_pose = agv_pose

        # tgt_trans = TransformStamped()
        # tgt_trans.header.frame_id = self.real_support_points.pose.header.frame_id
        # tgt_trans.header.stamp = rospy.Time.now()

        # tgt_trans.child_frame_id = self.real_base_link
        # tgt_trans.transform.translation.x = rot3[0,3]
        # tgt_trans.transform.translation.y = rot3[1,3]
        # tgt_trans.transform.translation.z = rot3[2,3]
        # quat = tf.transformations.quaternion_from_matrix(rot3)
        # tgt_trans.transform.rotation.x = quat[0]
        # tgt_trans.transform.rotation.y = quat[1]
        # tgt_trans.transform.rotation.z = quat[2]
        # tgt_trans.transform.rotation.w = quat[3]

        # self.tfBroadcaster.sendTransform(tgt_trans)

        # self.agv_pose_control()
        self.agv_joint_pub.publish(sent_command)

    def agv_pose_control_step(self):
        tgt_trans = TransformStamped()
        tgt_trans.header.frame_id = self.real_agv_pose.header.frame_id
        tgt_trans.header.stamp = rospy.Time.now()
        tgt_trans.child_frame_id = self.real_base_link

        tgt_trans.transform.translation.x = self.real_agv_pose.pose.position.x
        tgt_trans.transform.translation.y = self.real_agv_pose.pose.position.y
        tgt_trans.transform.translation.z = self.real_agv_hight

        tgt_trans.transform.rotation.x = self.real_agv_pose.pose.orientation.x
        tgt_trans.transform.rotation.y = self.real_agv_pose.pose.orientation.y
        tgt_trans.transform.rotation.z = self.real_agv_pose.pose.orientation.z
        tgt_trans.transform.rotation.w = self.real_agv_pose.pose.orientation.w
        self.tfBroadcaster.sendTransform(tgt_trans)

        xhist = self.agv_pose_hist['xhist']
        yhist = self.agv_pose_hist['yhist']
        zhist = self.agv_pose_hist['zhist']

        MAX_INT_LENGTH = self.MAX_INT_LENGTH 
        PID_I = self.PID_I 
        PID_P = self.PID_P 

        if (not self.tfBuffer.can_transform(self.base_link, self.real_base_link, rospy.Time(0))):
            return
        controlTransform = self.tfBuffer.lookup_transform(self.base_link, self.real_base_link, rospy.Time(0))

        quat = controlTransform.transform.rotation
        euler = tf.transformations.euler_from_quaternion([quat.x, quat.y, quat.z, quat.w],'sxyz')
        deltap = controlTransform.transform.translation
        # print(self.base_link,"x, y, theta:",deltap,euler)

        xhist.append(deltap.x)
        yhist.append(deltap.y)
        zhist.append(euler[2])

        xint = 0.0
        yint = 0.0
        zint = 0.0
        if(len(xhist)>MAX_INT_LENGTH):
            xhist.pop(0)
            yhist.pop(0)
            zhist.pop(0)

            xint = sum(xhist) * PID_I
            yint = sum(yhist) * PID_I
            zint = sum(zhist) * PID_I

        xint = xint + deltap.x * PID_P
        yint = yint + deltap.y * PID_P
        zint = zint + euler[2] * PID_P * 0.2 * 50

        cmd = Twist()
        cmd.linear.x = xint if math.fabs(xint) < 1.0 else math.fabs(xint)/xint * 1.0
        cmd.linear.y = yint if math.fabs(yint) < 1.0 else math.fabs(yint)/yint * 1.0
        cmd.angular.z = zint if math.fabs(zint) < 0.2*50 else math.fabs(zint)/zint * 0.2*50

        self.sim_cmd_vel.publish(cmd)

    def reset_agv_pose_control(self):
        self.agv_pose_hist={}
        self.agv_pose_hist['xhist'] = []
        self.agv_pose_hist['yhist'] = []
        self.agv_pose_hist['zhist'] = []

        self.MAX_INT_LENGTH = 10
        self.PID_I = 0.02
        self.PID_P = 1.0

    def run(self):
        rate = rospy.Rate(20)

        while not rospy.is_shutdown():
            # firstly get current state of the simulating model
            if (not self.tfBuffer.can_transform("world", self.base_link, rospy.Time(0))):
                rate.sleep()
                continue
            initPose = self.tfBuffer.lookup_transform('world', self.base_link, rospy.Time(0))
            
            curr_trans = PoseStamped()
            curr_trans.header.stamp = rospy.Time.now()

            curr_trans.pose.position.x = initPose.transform.translation.x
            curr_trans.pose.position.y = initPose.transform.translation.y
            curr_trans.pose.position.z = initPose.transform.translation.z

            curr_trans.pose.orientation.x = initPose.transform.rotation.x
            curr_trans.pose.orientation.y = initPose.transform.rotation.y
            curr_trans.pose.orientation.z = initPose.transform.rotation.z
            curr_trans.pose.orientation.w = initPose.transform.rotation.w

            # self.sim_agv_pose.publish(trans)
            self.real_agv_hight = curr_trans.pose.position.z

            self.lock.acquire()
            if (not self.real_agv_pose is None):
                self.agv_pose_control_step()
            # elif(not self.real_support_points is None):
            #     self.support_points_control_step()
            self.lock.release()
            
            rate.sleep()

if __name__ == '__main__':
    prefix = "pogo_agv"

    if(len(sys.argv)>1):
        prefix = sys.argv[1]
        print("file written to prefix",prefix)

    server = pogo_agv_service('agv_service',prefix)
    server.run()
